# -*- coding: utf-8 -*-
"""
Created on Wed Jan  5 21:11:28 2022

@author: 04566
"""


import socket
import numpy as np
import numba as nb
import open3d as o3d
import time

"""
ones_rslidar函数用于读取一个激光雷达数据集，该数据集包括x*10次雷达发回数据。
"""
@nb.jit()
def ones_rslidar(x,udp_socket):
    num=x
    temp_data=np.zeros(100)
    temp_data=temp_data[np.newaxis,:]
    for k in range(int(num*10)):
        (recv_data,addr) = udp_socket.recvfrom(1290)                  #利用UDP接受雷达数据
        temp=np.array(list(recv_data))                                #数据转存进numpy
        temp_data=np.vstack((temp_data,temp[42:1242].reshape(12,100)))#RSlidar一次的数据包包含12组有效数据，每组包含100个16进制数据
    alpha=np.add(temp_data[:,2]*256,temp_data[:,3])/100*3.1416/180    #每组第2个和第3个（从0开始）数据组成雷达水平旋转角度
    dises=np.zeros([int(temp_data.shape[0]),16])
    intenses=np.zeros([int(temp_data.shape[0]),16])
    for i in range(16):
        dises[:,i]=(temp_data[:,i*3+4]*256+temp_data[:,i*3+5])/100    #每组从第4个开始每3个数据为一个小组，每个小组中前2个数据代表距离
        intenses[:,i]=temp_data[:,i*3+6]                              #每组从第4个开始每3个数据为一个小组，每个小组中最后1个数据代表反射率
    return (dises,intenses,alpha)


@nb.jit()
def dis(dis_alpha,dis_sphere_angle,dis_dises):                        #利用水平角、俯仰角和距离计算坐标
    x=np.sin(dis_alpha)@np.cos(dis_sphere_angle)
    x=np.multiply(x,dis_dises)
    y=np.cos(dis_alpha)@np.cos(dis_sphere_angle)
    y=np.multiply(y,dis_dises)
    z=np.ones_like(x)*np.sin(dis_sphere_angle)
    z=np.multiply(z,dis_dises)
    x=x.reshape(-1,1)
    y=y.reshape(-1,1)
    z=z.reshape(-1,1)
    points=np.hstack((x,y,z))
    return points


vis = o3d.visualization.Visualizer()
vis.create_window( window_name='Open3D', width=960, height=540)
render_option: o3d.visualization.RenderOption = vis.get_render_option()	#设置点云渲染参数
render_option.point_size = 1.0	                                        #设置渲染点的大小
global_point_cloud=o3d.geometry.PointCloud()
last_point_cloud=o3d.geometry.PointCloud()
point_cloud = o3d.geometry.PointCloud()
point_cloud_np=np.zeros(3,)
point_cloud_np=point_cloud_np[np.newaxis,:]
trans = np.asarray([[1,0,0,0],                                          # 4x4 identity matrix，这是一个转换矩阵，
                    [0,1,0,0],                                          # 象征着没有任何位移，没有任何旋转，我们输入
                    [0,0,1,0],                                          # 这个矩阵为初始变换
                    [0,0,0,1]])
print(point_cloud_np.shape)
vis.add_geometry(global_point_cloud)
udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)           #建立UDP连接
udp_socket.bind(('192.168.1.102', 6699))
#channel=np.arange(8)                                                    #计算每个激光头俯仰角，或者俯仰角厂家提供
#sphere_angle=(channel*2-15)*3.1416/180
#sphere_angle=np.hstack((sphere_angle,(15-channel*2)*3.1416/180))
sphere_angle=np.array([-14.9927,-12.9868,-10.9859,-8.9785,\
                       -6.9945,-5.0006,-2.9838,-0.9918,\
                       15.0061,12.9936,11.0066,9.0065,\
                       7.0016,5.0006,2.9981,1.0026])
sphere_angle=sphere_angle*3.1416/180
sphere_angle=sphere_angle[np.newaxis,:]
to_reset = True
num=0
while True:
    
    (dises,intenses,alpha)=ones_rslidar(10,udp_socket)
    alpha=alpha[:,np.newaxis]
    dises=np.where(dises>600,0,dises)
    points=dis(alpha[1:],sphere_angle,dises[1:])
    point_cloud.points = o3d.utility.Vector3dVector(points)
    #temp_point_cloud, outlier_index = o3d.geometry.PointCloud.remove_radius_outlier(point_cloud,
                                          #nb_points=16,radius=0.5)
                                         
    threshold = 1.0  #移动范围的阀值
    trans_init = np.asarray([[1,0,0,0],   # 4x4 identity matrix，这是一个转换矩阵，
                             [0,1,0,0],   # 象征着没有任何位移，没有任何旋转，我们输入
                             [0,0,1,0],   # 这个矩阵为初始变换
                             [0,0,0,1]])
    
    point_cloud_base=[]
    point_cloud_imu=[]
    #运行icp
    if num==0:
        global_point_cloud.points=point_cloud.points
        last_point_cloud.points=point_cloud.points
                
    elif num>0 and num%2==0:               #icp计算激光雷达里程计
        reg_p2p = o3d.pipelines.registration.registration_icp(
                point_cloud, last_point_cloud, threshold, trans_init,
                o3d.pipelines.registration.TransformationEstimationPointToPoint(),
                o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=30))
    
        #print(reg_p2p)
        #icp_x,icp_y,icp_z=reg_p2p.transformation[:3,3]
        #if icp_x*icp_x+icp_y*icp_y+icp_z*icp_z>0.01:
        trans=trans@reg_p2p.transformation
        point_cloud_base.append(point_cloud)
        point_cloud_imu.append(trans)
        last_point_cloud.points=point_cloud.points
        print('small')
    else:
        pass
    if num>0 and num%50==0:                #icp误差修正，这部分以后还要改
        point_cloud.transform(trans)
        reg_p2p = o3d.pipelines.registration.registration_icp(
                point_cloud, global_point_cloud, threshold, trans_init,
                o3d.pipelines.registration.TransformationEstimationPointToPoint(),
                o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=100))
        delta=trans[:3,:]-reg_p2p.transformation[:3,:]
        #delta=np.zeros([3,4])
        for i in range(len(point_cloud_base)):
            new_imu=point_cloud_imu[i][:3,:]-delta*(1/len(point_cloud_base))*(i+1)
            new_imu=np.vstack((new_imu,np.array([[0,0,0,1]])))
            point_cloud_base[i].transform(new_imu)
            point_cloud_np=np.vstack((point_cloud_np,np.asarray(point_cloud_base[i].points)))
            temp_cloud=o3d.geometry.PointCloud()
            temp_cloud.points = o3d.utility.Vector3dVector(point_cloud_np)
            downpcd = temp_cloud.voxel_down_sample(voxel_size=0.1)
            global_point_cloud.points=downpcd.points
        point_cloud_base=[]
        point_cloud_imu=[]
        print('big')
    vis.update_geometry(global_point_cloud)
    if to_reset:
        vis.reset_view_point(True)
        to_reset = False
    vis.poll_events()
    vis.update_renderer()
    num+=1
    print(num)
    #time.sleep(0.1)
    #print(time.time()-thistime)
udp_socket.close()
vis.destroy_window()
